package site.wanjiahao.union;

import java.util.HashMap;
import java.util.Map;
import java.util.Objects;

/**
 * 自定义类型并查集
 */
public class GenericUnionFind<E> {

    // 存放并查集数据
    Map<E, Node<E>> containers = new HashMap<>();

    public E find(E e) {
        Node<E> node = findNode(e);
        return node == null? null: node.ele;
    }

    public void makeSet(E e) {
        containers.putIfAbsent(e, new Node<>(e));
    }

    public Node<E> findNode(E e) {
        Node<E> node = containers.get(e);
        if (node == null) return null;
        // 路径减半
        while (!Objects.equals(node.ele, node.parent.ele)) {
            node.parent = node.parent.parent;
            node = node.parent;
        }
        return node;
    }

    // 基于rank的优化
    public void union(E e1, E e2) {
        Node<E> node1 = findNode(e1);
        Node<E> node2 = findNode(e2);
        if (Objects.equals(node1, node2)) return;

        if (node1.rank > node2.rank) {
            node2.parent = node1;
        } else if (node1.rank < node2.rank) {
            node1.parent = node2;
        } else {
            node2.parent = node1;
            node1.rank++;
        }
    }

    public boolean isSame(E e1, E e2) {
        return find(e1) == find(e2);
    }


   private static class Node<E> {
       E ele;

       Node<E> parent = this;

       int rank = 1;

       public Node(E ele) {
           this.ele = ele;
       }
   }
}
